# IDS 过滤
from time import time_ns

from scapy.compat import raw
from scapy.layers.inet import IP, TCP

from NFs.ids_rule import IDSRule


class IDS(object):
    def __init__(self, file='ids_rules.csv'):
        self.file = file
        self.rules = self.pares_rules()

    def pares_rules(self):
        rules = list()
        with open(self.file, 'r') as f:
            for line in f:
                rules.append(IDSRule(line))
        return rules

    def match(self, pkt: bytes) -> bool:
        scapy_pkt = IP(pkt)
        for rule in self.rules:
            res = rule.match(scapy_pkt)
            if res:
                return True
        return False


if __name__ == '__main__':
    ids = IDS()
    print(len(ids.rules))
    start = time_ns()
    pkt = b'E\x00\x00@\x00\x01\x00\x00@\x06\x11\xe1\xf7\x89\xe9\xcc$\x91b\xeffeM\xf1\x00\x00\x00\x00\x00\x00\x00\x00P' \
          b'\x02 \x00\xe2\x0c\x00\x00aaaaaaaaaaaaaaaaaaaaaaaa '
    count = 1
    pkt = IP(src='247.137.233.204', dst='36.145.98.239') / TCP(sport=26213, dport=19953)
    pkt[TCP].add_payload(b'a' * 3996)
    pkt = raw(pkt)
    for i in range(100000):
        ids.match(pkt)
    print((time_ns() - start) / 100000 / 1000)  # 平均每条指令执行时间
